﻿///////////////////////////////////////////////////////////////////////////////
//
//  This file is part of MathLib.NET.
//
//  This library is free software; you can redistribute it and/or
//  modify it under the terms of the GNU Lesser General Public
//  License as published by the Free Software Foundation; either
//  version 2.1 of the License, or (at your option) any later version.
//  
//  This library is distributed in the hope that it will be useful,
//  but WITHOUT ANY WARRANTY; without even the implied warranty of
//  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
//  Lesser General Public License for more details.
//  
//  You should have received a copy of the GNU Lesser General Public
//  License along with this library;  If not, see 
//  <http://www.gnu.org/licenses/>.
//
///////////////////////////////////////////////////////////////////////////////

using System;
using MathLib;
using ILNumerics;
using NUnit.Framework;

namespace MathLibTest
{
    [TestFixture]
    public class FourierTest
    {
        /// <summary>
        /// This test computes 127 256-point FFTs in a single call to fft().
        /// 
        /// The input to each transform is a cosine wave of N wavelengths
        /// where N is the number of the transform (0..126).
        /// 
        /// The transformation is performed along dimension 0.
        /// 
        /// The output is tested to ensure that the peak frequency cell is
        /// equal to N.
        /// </summary>
        [Test]
        public void fft_Ex1()
        {
            int FFT_SIZE = 256;

            ILArray<double> input = new ILArray<double>(new int[] { FFT_SIZE, FFT_SIZE / 2 - 1 });

            // Generate 2d sine-waves
            for (int j = 0; j < FFT_SIZE / 2 - 1; j++)
                for (int i = 0; i < FFT_SIZE; i++)
                    input.InternalArray4Experts[j * FFT_SIZE + i] = Math.Cos((double)(i * j) * 2.0 * Math.PI / (double)(FFT_SIZE));

            // Transform
            ILArray<complex> output = MLMath.fft(input);

            // Test output
            for (int j = 0; j < FFT_SIZE / 2 - 1; j++)
            {
                // Find the cell containing the peak value
                double peakValue = 0.0;
                int peakCell = 0;
                for (int i = 0; i < FFT_SIZE / 2; i++)
                    if (complex.Abs(output.InternalArray4Experts[j * FFT_SIZE + i]) >= peakValue)
                    {
                        peakValue = complex.Abs(output.InternalArray4Experts[j * FFT_SIZE + i]);
                        peakCell = i;
                    }

                // Verify that the peak is in the correct cell
                Assert.IsTrue(peakCell == j);
                //Console.WriteLine("{0} = {1} ({2})", j, peakCell, Math.Abs(j - peakCell));
            }
        }

        /// <summary>
        /// This test computes 127 256-point FFTs in a single call to fft().
        /// 
        /// The input to each transform is a cosine wave of N wavelengths
        /// where N is the number of the transform (0..126).
        /// 
        /// The transformation is performed along dimension 1.
        /// 
        /// The output is tested to ensure that the peak frequency cell is
        /// equal to N.
        /// </summary>
        [Test]
        public void fft_Ex2()
        {
            int FFT_SIZE = 256;

            ILArray<double> input = new ILArray<double>(new int[] { FFT_SIZE / 2 - 1, FFT_SIZE });

            // Generate 2d sine-waves
            for (int j = 0; j < FFT_SIZE / 2 - 1; j++)
                for (int i = 0; i < FFT_SIZE; i++)
                    input.InternalArray4Experts[i * (FFT_SIZE / 2 - 1) + j] = Math.Cos((double)(i * j) * 2.0 * Math.PI / (double)(FFT_SIZE));

            // Transform
            ILArray<complex> output = MLMath.fft(input, null, 1);

            // Test output
            for (int j = 0; j < FFT_SIZE / 2 - 1; j++)
            {
                // Find the cell containing the peak value
                double peakValue = 0.0;
                int peakCell = 0;
                for (int i = 0; i < FFT_SIZE / 2; i++)
                    if (complex.Abs(output.InternalArray4Experts[i * (FFT_SIZE / 2 - 1) + j]) >= peakValue)
                    {
                        peakValue = complex.Abs(output.InternalArray4Experts[i * (FFT_SIZE / 2 - 1) + j]);
                        peakCell = i;
                    }

                // Verify that the peak is in the correct cell
                Assert.IsTrue(peakCell == j);
                //Console.WriteLine("{0} = {1} ({2})", j, peakCell, Math.Abs(j - peakCell));
            }
        }

        /// <summary>
        /// This test computes 64 256-point FFTs in a single call to fft().
        /// 
        /// The input to each transform is a cosine wave of N wavelengths
        /// where N is the number of the transform (0..63).
        /// 
        /// The transformation is performed along dimension 1 and the 64
        /// transforms are distributed across two dimensions yielding a
        /// input dimensionality of [8][256][8]
        /// 
        /// The output is tested to ensure that the peak frequency cell is
        /// equal to N.
        /// </summary>
        [Test]
        public void fft_Ex3()
        {
            int FFT_SIZE = 256;
            int FFTS_PER_DIM = 8;

            ILArray<double> input = new ILArray<double>(new int[] { FFTS_PER_DIM, FFT_SIZE, FFTS_PER_DIM });

            // Generate 3d sine-waves
            for (int k = 0; k < FFTS_PER_DIM; k++)
                for (int j = 0; j < FFT_SIZE; j++)
                    for (int i = 0; i < FFTS_PER_DIM; i++)
                        input.InternalArray4Experts[(k * FFT_SIZE + j) * FFTS_PER_DIM + i] = Math.Cos((double)((k * FFTS_PER_DIM + i) * j) * 2.0 * Math.PI / (double)(FFT_SIZE));

            // Transform
            ILArray<complex> output = MLMath.fft(input, null, 1);

            // Test output
            for (int k = 0; k < FFTS_PER_DIM; k++)
            {
                for (int i = 0; i < FFTS_PER_DIM; i++)
                {
                    // Find the cell containing the peak value
                    double peakValue = 0.0;
                    int peakCell = 0;
                    for (int j = 0; j < FFT_SIZE / 2; j++)
                        if (complex.Abs(output.InternalArray4Experts[(k * FFT_SIZE + j) * FFTS_PER_DIM + i]) >= peakValue)
                        {
                            peakValue = complex.Abs(output.InternalArray4Experts[(k * FFT_SIZE + j) * FFTS_PER_DIM + i]);
                            peakCell = j;
                        }

                    // Verify that the peak is in the correct cell
                    Assert.IsTrue(peakCell == k * FFTS_PER_DIM + i);
                    //Console.WriteLine("{0},{1} = {2} ({3})", i, k, peakCell, Math.Abs(k * FFTS_PER_DIM + i - peakCell));
                }
            }
        }

        /// <summary>
        /// This test computes 64 128-point FFTs in a single call to fft().
        /// 
        /// The input to each transform is a 256-point cosine wave of 2 * N
        /// wavelengths where N is the number of the transform (0..63).
        /// The input is truncated to 128-points prior to transformation.
        /// 
        /// The transformation is performed along dimension 1 and the 64
        /// transforms are distributed across two dimensions yielding a
        /// dimensionality of [8][128][8]
        /// 
        /// The output is tested to ensure that the peak frequency cell is
        /// equal to N.
        /// </summary>
        [Test]
        public void fft_Ex4()
        {
            int FFT_SIZE = 256;
            int FFTS_PER_DIM = 8;

            ILArray<double> input = new ILArray<double>(new int[] { FFTS_PER_DIM, FFT_SIZE, FFTS_PER_DIM });

            // Generate 3d sine-waves
            for (int k = 0; k < FFTS_PER_DIM; k++)
                for (int j = 0; j < FFT_SIZE; j++)
                    for (int i = 0; i < FFTS_PER_DIM; i++)
                        input.InternalArray4Experts[(k * FFT_SIZE + j) * FFTS_PER_DIM + i] = Math.Cos((double)((k * FFTS_PER_DIM + i) * j) * 4.0 * Math.PI / (double)(FFT_SIZE));

            // Transform
            ILArray<complex> output = MLMath.fft(input, FFT_SIZE / 2, 1);

            // Test output
            for (int k = 0; k < FFTS_PER_DIM; k++)
            {
                for (int i = 0; i < FFTS_PER_DIM; i++)
                {
                    // Find the cell containing the peak value
                    double peakValue = 0.0;
                    int peakCell = 0;
                    for (int j = 0; j < FFT_SIZE / 4; j++)
                        if (complex.Abs(output.InternalArray4Experts[(k * FFT_SIZE / 2 + j) * FFTS_PER_DIM + i]) >= peakValue)
                        {
                            peakValue = complex.Abs(output.InternalArray4Experts[(k * FFT_SIZE / 2 + j) * FFTS_PER_DIM + i]);
                            peakCell = j;
                        }

                    // Verify that the peak is in the correct cell
                    Assert.IsTrue(peakCell == k * FFTS_PER_DIM + i);
                    //Console.WriteLine("{0},{1} = {2} ({3})", i, k, peakCell, Math.Abs(k * FFTS_PER_DIM + i - peakCell));
                }
            }
        }

        /// <summary>
        /// This test computes 127 256-point complex FFTs in a single call to fft().
        /// 
        /// The input to each transform is a cosine wave of N wavelengths
        /// where N is the number of the transform (0..126).
        /// 
        /// The transformation is performed along dimension 0.
        /// 
        /// The output is tested to ensure that the peak frequency cell is
        /// equal to N.
        /// </summary>
        [Test]
        public void fft_Ex5()
        {
            int FFT_SIZE = 256;

            ILArray<complex> input = new ILArray<complex>(new int[] { FFT_SIZE, FFT_SIZE / 2 - 1 });

            // Generate 2d sine-waves
            for (int j = 0; j < FFT_SIZE / 2 - 1; j++)
                for (int i = 0; i < FFT_SIZE; i++)
                    input.InternalArray4Experts[j * FFT_SIZE + i] = Math.Cos((double)(i * j) * 2.0 * Math.PI / (double)(FFT_SIZE));

            // Transform
            ILArray<complex> output = MLMath.fft(input);

            // Test output
            for (int j = 0; j < FFT_SIZE / 2 - 1; j++)
            {
                // Find the cell containing the peak value
                double peakValue = 0.0;
                int peakCell = 0;
                for (int i = 0; i < FFT_SIZE / 2; i++)
                    if (complex.Abs(output.InternalArray4Experts[j * FFT_SIZE + i]) >= peakValue)
                    {
                        peakValue = complex.Abs(output.InternalArray4Experts[j * FFT_SIZE + i]);
                        peakCell = i;
                    }

                // Verify that the peak is in the correct cell
                Assert.IsTrue(peakCell == j);
                //Console.WriteLine("{0} = {1} ({2})", j, peakCell, Math.Abs(j - peakCell));
            }
        }

        /// <summary>
        /// This test computes 127 256-point complex FFTs in a single call to fft().
        /// 
        /// The input to each transform is a cosine wave of N wavelengths
        /// where N is the number of the transform (0..126).
        /// 
        /// The transformation is performed along dimension 1.
        /// 
        /// The output is tested to ensure that the peak frequency cell is
        /// equal to N.
        /// </summary>
        [Test]
        public void fft_Ex6()
        {
            int FFT_SIZE = 256;

            ILArray<complex> input = new ILArray<complex>(new int[] { FFT_SIZE / 2 - 1, FFT_SIZE });

            // Generate 2d sine-waves
            for (int j = 0; j < FFT_SIZE / 2 - 1; j++)
                for (int i = 0; i < FFT_SIZE; i++)
                    input.InternalArray4Experts[i * (FFT_SIZE / 2 - 1) + j] = Math.Cos((double)(i * j) * 2.0 * Math.PI / (double)(FFT_SIZE));

            // Transform
            ILArray<complex> output = MLMath.fft(input, null, 1);

            // Test output
            for (int j = 0; j < FFT_SIZE / 2 - 1; j++)
            {
                // Find the cell containing the peak value
                double peakValue = 0.0;
                int peakCell = 0;
                for (int i = 0; i < FFT_SIZE / 2; i++)
                    if (complex.Abs(output.InternalArray4Experts[i * (FFT_SIZE / 2 - 1) + j]) >= peakValue)
                    {
                        peakValue = complex.Abs(output.InternalArray4Experts[i * (FFT_SIZE / 2 - 1) + j]);
                        peakCell = i;
                    }

                // Verify that the peak is in the correct cell
                Assert.IsTrue(peakCell == j);
                //Console.WriteLine("{0} = {1} ({2})", j, peakCell, Math.Abs(j - peakCell));
            }
        }

        /// <summary>
        /// This test computes 64 256-point complex FFTs in a single call to fft().
        /// 
        /// The input to each transform is a cosine wave of N wavelengths
        /// where N is the number of the transform (0..63).
        /// 
        /// The transformation is performed along dimension 1 and the 64
        /// transforms are distributed across two dimensions yielding a
        /// input dimensionality of [8][256][8]
        /// 
        /// The output is tested to ensure that the peak frequency cell is
        /// equal to N.
        /// </summary>
        [Test]
        public void fft_Ex7()
        {
            int FFT_SIZE = 256;
            int FFTS_PER_DIM = 8;

            ILArray<complex> input = new ILArray<complex>(new int[] { FFTS_PER_DIM, FFT_SIZE, FFTS_PER_DIM });

            // Generate 3d sine-waves
            for (int k = 0; k < FFTS_PER_DIM; k++)
                for (int j = 0; j < FFT_SIZE; j++)
                    for (int i = 0; i < FFTS_PER_DIM; i++)
                        input.InternalArray4Experts[(k * FFT_SIZE + j) * FFTS_PER_DIM + i] = Math.Cos((double)((k * FFTS_PER_DIM + i) * j) * 2.0 * Math.PI / (double)(FFT_SIZE));

            // Transform
            ILArray<complex> output = MLMath.fft(input, null, 1);

            // Test output
            for (int k = 0; k < FFTS_PER_DIM; k++)
            {
                for (int i = 0; i < FFTS_PER_DIM; i++)
                {
                    // Find the cell containing the peak value
                    double peakValue = 0.0;
                    int peakCell = 0;
                    for (int j = 0; j < FFT_SIZE / 2; j++)
                        if (complex.Abs(output.InternalArray4Experts[(k * FFT_SIZE + j) * FFTS_PER_DIM + i]) >= peakValue)
                        {
                            peakValue = complex.Abs(output.InternalArray4Experts[(k * FFT_SIZE + j) * FFTS_PER_DIM + i]);
                            peakCell = j;
                        }

                    // Verify that the peak is in the correct cell
                    Assert.IsTrue(peakCell == k * FFTS_PER_DIM + i);
                    //Console.WriteLine("{0},{1} = {2} ({3})", i, k, peakCell, Math.Abs(k * FFTS_PER_DIM + i - peakCell));
                }
            }
        }

        /// <summary>
        /// This test computes 64 128-point complex FFTs in a single call to fft().
        /// 
        /// The input to each transform is a 256-point cosine wave of 2 * N
        /// wavelengths where N is the number of the transform (0..63).
        /// The input is truncated to 128-points prior to transformation.
        /// 
        /// The transformation is performed along dimension 1 and the 64
        /// transforms are distributed across two dimensions yielding a
        /// dimensionality of [8][128][8]
        /// 
        /// The output is tested to ensure that the peak frequency cell is
        /// equal to N.
        /// </summary>
        [Test]
        public void fft_Ex8()
        {
            int FFT_SIZE = 256;
            int FFTS_PER_DIM = 8;

            ILArray<complex> input = new ILArray<complex>(new int[] { FFTS_PER_DIM, FFT_SIZE, FFTS_PER_DIM });

            // Generate 3d sine-waves
            for (int k = 0; k < FFTS_PER_DIM; k++)
                for (int j = 0; j < FFT_SIZE; j++)
                    for (int i = 0; i < FFTS_PER_DIM; i++)
                        input.InternalArray4Experts[(k * FFT_SIZE + j) * FFTS_PER_DIM + i] = Math.Cos((double)((k * FFTS_PER_DIM + i) * j) * 4.0 * Math.PI / (double)(FFT_SIZE));

            // Transform
            ILArray<complex> output = MLMath.fft(input, FFT_SIZE / 2, 1);

            // Test output
            for (int k = 0; k < FFTS_PER_DIM; k++)
            {
                for (int i = 0; i < FFTS_PER_DIM; i++)
                {
                    // Find the cell containing the peak value
                    double peakValue = 0.0;
                    int peakCell = 0;
                    for (int j = 0; j < FFT_SIZE / 4; j++)
                        if (complex.Abs(output.InternalArray4Experts[(k * FFT_SIZE / 2 + j) * FFTS_PER_DIM + i]) >= peakValue)
                        {
                            peakValue = complex.Abs(output.InternalArray4Experts[(k * FFT_SIZE / 2 + j) * FFTS_PER_DIM + i]);
                            peakCell = j;
                        }

                    // Verify that the peak is in the correct cell
                    Assert.IsTrue(peakCell == k * FFTS_PER_DIM + i);
                    //Console.WriteLine("{0},{1} = {2} ({3})", i, k, peakCell, Math.Abs(k * FFTS_PER_DIM + i - peakCell));
                }
            }
        }

        /// <summary>
        /// This test computes 64 256-point complex FFTs in a single call to fft()
        /// and then transforms the result back into the time-domain with a single
        /// call to ifft().
        /// 
        /// The input to each transform is a cosine wave of N wavelengths
        /// where N is the number of the transform (0..63).
        /// 
        /// The transformation is performed along dimension 1 and the 64
        /// transforms are distributed across two dimensions yielding a
        /// input dimensionality of [8][256][8].
        /// 
        /// The output is tested to ensure that it is the same as the original input.
        /// </summary>
        [Test]
        public void ifft_Ex1()
        {
            int FFT_SIZE = 256;
            int FFTS_PER_DIM = 8;
            double TOLERANCE = 0.0001;

            ILArray<complex> input = new ILArray<complex>(new int[] { FFTS_PER_DIM, FFT_SIZE, FFTS_PER_DIM });

            // Generate 3d sine-waves
            for (int k = 0; k < FFTS_PER_DIM; k++)
                for (int j = 0; j < FFT_SIZE; j++)
                    for (int i = 0; i < FFTS_PER_DIM; i++)
                        input.InternalArray4Experts[(k * FFT_SIZE + j) * FFTS_PER_DIM + i] = Math.Cos((double)((k * FFTS_PER_DIM + i) * j) * 2.0 * Math.PI / (double)(FFT_SIZE));

            // Transform
            ILArray<complex> output = MLMath.ifft(MLMath.fft(input, null, 1), null, 1);

            // Test output
            for (int k = 0; k < FFTS_PER_DIM; k++)
                for (int j = 0; j < FFT_SIZE; j++)
                    for (int i = 0; i < FFTS_PER_DIM; i++)
                        Assert.IsTrue(Math.Abs(output.InternalArray4Experts[(k * FFT_SIZE + j) * FFTS_PER_DIM + i].real - Math.Cos((double)((k * FFTS_PER_DIM + i) * j) * 2.0 * Math.PI / (double)(FFT_SIZE))) < TOLERANCE);
        }

        /// <summary>
        /// This test computes 16 256x256 2D FFT transformations for an input
        /// signal consisting of cross-multiplied cosine waves.
        /// 
        /// The output is tested to ensure that the peak cell corresponds to
        /// the frequency of the input signal.
        /// </summary>
        [Test]
        public void fft2_Ex1()
        {
            int FFT_SIZE = 256;
            int NUM_TEST_FREQ = 16;

            ILArray<double> input = new ILArray<double>(new int[] { FFT_SIZE, FFT_SIZE });

            for (int f = 0; f < FFT_SIZE / 2; f += FFT_SIZE / (2 * NUM_TEST_FREQ))
            {
                // Generate 2d sinusoid
                for (int i = 0; i < FFT_SIZE; i++)
                    for (int j = 0; j < FFT_SIZE; j++)
                        input.InternalArray4Experts[j * FFT_SIZE + i] = Math.Cos((double)(i * f) * 2.0 * Math.PI / (double)(FFT_SIZE)) * Math.Cos((double)(j * f) * 2.0 * Math.PI / (double)(FFT_SIZE));

                // Transform
                ILArray<complex> output = MLMath.fft2(input);

                // Find the cell containing the peak value
                double peakValue = 0.0;
                int peakCell = 0;
                for (int i = 0; i < FFT_SIZE / 2; i++)
                    for (int j = 0; j < FFT_SIZE / 2; j++)
                        if (complex.Abs(output.InternalArray4Experts[j * FFT_SIZE + i]) >= peakValue)
                        {
                            peakValue = complex.Abs(output.InternalArray4Experts[j * FFT_SIZE + i]);
                            peakCell = j * FFT_SIZE + i;
                        }

                // Calculate column and row of peak cell
                int peakCol = peakCell % FFT_SIZE;
                int peakRow = peakCell / FFT_SIZE;

                // Verify that the peak is in the correct cell
                Assert.IsTrue(peakCol == f && peakRow == f);
            }
        }
    }
}
